import re
from datetime import datetime
import httpx
import json
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP
import mcp.types as types

import os
# import logging
# LOG_FILE = "D:/scripts/python/myfastapi/logs/ticketServer.log"
# os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)
# # 配置日志时指定文件编码为 UTF-8
# logging.basicConfig(
#     level=logging.DEBUG,
#     format="%(asctime)s - %(levelname)s - %(message)s",
#     handlers=[
#         logging.FileHandler(LOG_FILE, encoding="utf-8"),  # 写入文件时指定编码
#         logging.StreamHandler()  # 控制台输出（依赖终端编码）
#     ]
# )
# logger = logging.getLogger(__name__)
# # logger.info(f"🔍 查询参数: {} → {} ({})")

# 创建MCP服务器实例
mcp = FastMCP(name="variflight-mcp")

load_dotenv(override=True)
api_key = os.getenv('VARIFLIGHT_API_KEY')
api_url = "https://mcp.variflight.com/api/v1/mcp/data"

class OpenAlService:
    def __init__(self):
        self.base_url = api_url
        self.api_key = api_key

    async def make_request(self, endpoint, params):
        request_body = {
            'endpoint': endpoint,
            'params': params
        }

        headers = {
            'X-VARIFLIGHT-KEY': self.api_key or '',
            'Content-Type': 'application/json',
        }

        try:
            async with httpx.AsyncClient() as client:
                response = await client.post(self.base_url, headers=headers, json=request_body)
                response.raise_for_status()
                return response.json()
        except httpx.HTTPError as e:
            raise Exception(f'API request failed: {str(e)}')

    async def get_flights_by_dep_arr(self, dep=None, depcity=None, arr=None, arrcity=None, date=None):
        return await self.make_request('flights', {
            'dep': dep,
            'depcity': depcity,
            'arr': arr,
            'arrcity': arrcity,
            'date': date,
        })

    async def get_flight_by_number(self, fnum, date, dep=None, arr=None):
        params = {
            'fnum': fnum,
            'date': date,
        }
        if dep:
            params['dep'] = dep
        if arr:
            params['arr'] = arr
        return await self.make_request('flight', params)

    async def get_flight_transfer_info(self, depcity, arrcity, depdate):
        return await self.make_request('transfer', {
            'depcity': depcity,
            'arrcity': arrcity,
            'depdate': depdate
        })

    async def get_realtime_location_by_anum(self, anum):
        return await self.make_request('realtimeLocation', {
            'anum': anum
        })

    async def get_airport_weather(self, airport):
        return await self.make_request('futureAirportWeather', {
            'code': airport,
            'type': '1'
        })

    async def get_flight_happiness_index(self, fnum, date, dep=None, arr=None):
        params = {
            'fnum': fnum,
            'date': date,
        }
        if dep:
            params['dep'] = dep
        if arr:
            params['arr'] = arr
        return await self.make_request('happiness', params)

    async def search_flight_itineraries(self, dep_city_code, arr_city_code, dep_date):
        return await self.make_request('searchFlightItineraries', {
            'depCityCode': dep_city_code,
            'arrCityCode': arr_city_code,
            'depDate': dep_date
        })


flight_service = OpenAlService()


# 参数验证函数
def validate_dep_arr_params(params):
    if params.get('dep') and not re.match(r'^[A-Z]{3}$', params['dep']):
        raise ValueError("Departure airport code must be 3 uppercase letters")
    if params.get('depcity') and not re.match(r'^[A-Z]{3}$', params['depcity']):
        raise ValueError("Departure city code must be 3 uppercase letters")
    if params.get('arr') and not re.match(r'^[A-Z]{3}$', params['arr']):
        raise ValueError("Arrival airport code must be 3 uppercase letters")
    if params.get('arrcity') and not re.match(r'^[A-Z]{3}$', params['arrcity']):
        raise ValueError("Arrival city code must be 3 uppercase letters")
    if params.get('date') and not re.match(r'^\d{4}-\d{2}-\d{2}$', params['date']):
        raise ValueError("Date must be in YYYY-MM-DD format")


def validate_flight_number_params(params):
    if not re.match(r'^[A-Z0-9]{2,3}[0-9]{1,4}$', params.get('fnum', '')):
        raise ValueError("Flight number format is invalid")
    if params.get('date') and not re.match(r'^\d{4}-\d{2}-\d{2}$', params.get('date')):
        raise ValueError("Date must be in YYYY-MM-DD format")
    if params.get('dep') and not re.match(r'^[A-Z]{3}$', params.get('dep')):
        raise ValueError("Departure airport code must be 3 uppercase letters")
    if params.get('arr') and not re.match(r'^[A-Z]{3}$', params.get('arr')):
        raise ValueError("Arrival airport code must be 3 uppercase letters")


def validate_transfer_params(params):
    if not re.match(r'^\d{4}-\d{2}-\d{2}$', params.get('depdate')):
        raise ValueError("Date must be in YYYY-MM-DD format")
    if not re.match(r'^[A-Z]{3}$', params.get('depcity')):
        raise ValueError("Departure city code must be 3 uppercase letters")
    if not re.match(r'^[A-Z]{3}$', params.get('arrcity')):
        raise ValueError("Arrival city code must be 3 uppercase letters")


def validate_realtime_location_params(params):
    if not params.get('anum'):
        raise ValueError("Aircraft number is required")


def validate_itinerary_params(params):
    if not re.match(r'^[A-Z]{3}$', params.get('depCityCode', '')):
        raise ValueError("Departure city code must be 3 uppercase letters")
    if not re.match(r'^[A-Z]{3}$', params.get('arrCityCode', '')):
        raise ValueError("Arrival city code must be 3 uppercase letters")
    if not re.match(r'^\d{4}-\d{2}-\d{2}$', params.get('depDate', '')):
        raise ValueError("Date must be in YYYY-MM-DD format")


# 工具函数
async def search_flights_by_dep_arr(
        name: str, arguments: dict
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """通过出发地和目的地查询航班"""
    try:
        validate_dep_arr_params(arguments)
        flights = await flight_service.get_flights_by_dep_arr(
            arguments.get('dep'),
            arguments.get('depcity'),
            arguments.get('arr'),
            arguments.get('arrcity'),
            arguments.get('date')
        )
        return [types.TextContent(type="text", text=json.dumps(flights, indent=2, ensure_ascii=False))]
    except Exception as e:
        raise Exception(f"Error searching flights by dep/arr: {str(e)}")


async def search_flights_by_number(
        name: str, arguments: dict
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """通过航班号查询航班"""
    try:
        validate_flight_number_params(arguments)
        flights = await flight_service.get_flight_by_number(
            arguments.get('fnum'),
            arguments.get('date'),
            arguments.get('dep'),
            arguments.get('arr')
        )
        return [types.TextContent(type="text", text=json.dumps(flights, indent=2, ensure_ascii=False))]
    except Exception as e:
        raise Exception(f"Error searching flights by number: {str(e)}")


async def get_flight_transfer_info(
        name: str, arguments: dict
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """获取航班中转信息"""
    try:
        validate_transfer_params(arguments)
        flights = await flight_service.get_flight_transfer_info(
            arguments.get('depcity'),
            arguments.get('arrcity'),
            arguments.get('depdate')
        )
        return [types.TextContent(type="text", text=json.dumps(flights, indent=2, ensure_ascii=False))]
    except Exception as e:
        raise Exception(f"Error getting flight transfer info: {str(e)}")


async def get_flight_happiness_index(
        name: str, arguments: dict
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """获取航班舒适度指数"""
    try:
        validate_flight_number_params(arguments)
        flights = await flight_service.get_flight_happiness_index(
            arguments.get('fnum'),
            arguments.get('date'),
            arguments.get('dep'),
            arguments.get('arr')
        )
        return [types.TextContent(type="text", text=json.dumps(flights, indent=2, ensure_ascii=False))]
    except Exception as e:
        raise Exception(f"Error getting flight happiness index: {str(e)}")


async def get_realtime_location_by_anum(
        name: str, arguments: dict
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """飞机实时位置查询"""
    try:
        validate_realtime_location_params(arguments)
        realtime_location = await flight_service.get_realtime_location_by_anum(
            arguments.get('anum')
        )
        return [types.TextContent(type="text", text=json.dumps(realtime_location, indent=2, ensure_ascii=False))]
    except Exception as e:
        raise Exception(f"Error getting realtime location: {str(e)}")


async def get_today_date(
        name: str, arguments: dict
) -> list:
    """只返回当前时间（YYYY-MM-DD），不返回相对日期、周几等。"""
    today = datetime.now()
    text = today.strftime('%Y-%m-%d')
    return [{"type": "text", "text": text}]


async def search_flight_itineraries(
        name: str, arguments: dict
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """搜索航班方案"""
    try:
        validate_itinerary_params(arguments)
        flight_itineraries = await flight_service.search_flight_itineraries(
            arguments.get('depCityCode'),
            arguments.get('arrCityCode'),
            arguments.get('depDate')
        )
        return [types.TextContent(type="text", text=json.dumps(flight_itineraries, indent=2, ensure_ascii=False))]
    except Exception as e:
        raise Exception(f"Error searching flight itineraries: {str(e)}")


async def list_tools() -> list[types.Tool]:
    """列出所有可用的工具"""
    return [
        types.Tool(
            name="searchFlightsByDepArr",
            description="通过出发地和目的地查询航班，支持按机场代码或城市代码搜索",
            inputSchema={
                "type": "object",
                "properties": {
                    "dep": {
                        "type": "string",
                        "description": "出发地机场三字码（如：PEK）"
                    },
                    "depcity": {
                        "type": "string",
                        "description": "出发地城市三字码"
                    },
                    "arr": {
                        "type": "string",
                        "description": "目的地机场三字码（如：SHA）"
                    },
                    "arrcity": {
                        "type": "string",
                        "description": "目的地城市三字码"
                    },
                    "date": {
                        "type": "string",
                        "description": "日期，格式为YYYY-MM-DD"
                    }
                }
            }
        ),
        types.Tool(
            name="searchFlightsByNumber",
            description="通过航班号查询航班信息",
            inputSchema={
                "type": "object",
                "required": ["fnum", "date"],
                "properties": {
                    "fnum": {
                        "type": "string",
                        "description": "航班号（如：CA1234）"
                    },
                    "date": {
                        "type": "string",
                        "description": "日期，格式为YYYY-MM-DD"
                    },
                    "dep": {
                        "type": "string",
                        "description": "出发地机场三字码（可选）"
                    },
                    "arr": {
                        "type": "string",
                        "description": "目的地机场三字码（可选）"
                    }
                }
            }
        ),
        types.Tool(
            name="getFlightTransferInfo",
            description="获取航班中转信息",
            inputSchema={
                "type": "object",
                "required": ["depcity", "arrcity", "depdate"],
                "properties": {
                    "depcity": {
                        "type": "string",
                        "description": "出发地城市三字码"
                    },
                    "arrcity": {
                        "type": "string",
                        "description": "目的地城市三字码"
                    },
                    "depdate": {
                        "type": "string",
                        "description": "出发日期，格式为YYYY-MM-DD"
                    }
                }
            }
        ),
        types.Tool(
            name="flightHappinessIndex",
            description="获取航班舒适度指数",
            inputSchema={
                "type": "object",
                "required": ["fnum", "date"],
                "properties": {
                    "fnum": {
                        "type": "string",
                        "description": "航班号"
                    },
                    "date": {
                        "type": "string",
                        "description": "日期，格式为YYYY-MM-DD"
                    },
                    "dep": {
                        "type": "string",
                        "description": "出发地机场三字码（可选）"
                    },
                    "arr": {
                        "type": "string",
                        "description": "目的地机场三字码（可选）"
                    }
                }
            }
        ),
        types.Tool(
            name="getRealtimeLocationByAnum",
            description="通过飞机注册号查询实时位置",
            inputSchema={
                "type": "object",
                "required": ["anum"],
                "properties": {
                    "anum": {
                        "type": "string",
                        "description": "飞机注册号"
                    }
                }
            }
        ),
        types.Tool(
            name="getTodayDate",
            description="获取当前日期信息，支持相对日期计算。",
            inputSchema={
                "type": "object",
                "title": "获取当前时间参数",
                "properties": {
                    "format": {"type": "string", "title": "日期格式", "description": "返回的日期格式，默认为YYYY-MM-DD",
                               "default": "YYYY-MM-DD"}
                },
                "additionalProperties": False
            }
        ),
        types.Tool(
            name="searchFlightItineraries",
            description="搜索可购买的航班方案",
            inputSchema={
                "type": "object",
                "required": ["depCityCode", "arrCityCode", "depDate"],
                "properties": {
                    "depCityCode": {
                        "type": "string",
                        "description": "出发城市三字码"
                    },
                    "arrCityCode": {
                        "type": "string",
                        "description": "到达城市三字码"
                    },
                    "depDate": {
                        "type": "string",
                        "description": "出发日期，格式为YYYY-MM-DD"
                    }
                }
            }
        )
    ]


async def dispatch(
        name: str, arguments: dict
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """分发工具调用"""
    match name:
        case "searchFlightsByDepArr":
            return await search_flights_by_dep_arr(name, arguments)
        case "searchFlightsByNumber":
            return await search_flights_by_number(name, arguments)
        case "getFlightTransferInfo":
            return await get_flight_transfer_info(name, arguments)
        case "flightHappinessIndex":
            return await get_flight_happiness_index(name, arguments)
        case "getRealtimeLocationByAnum":
            return await get_realtime_location_by_anum(name, arguments)
        case "getTodayDate":
            return await get_today_date(name, arguments)
        case "searchFlightItineraries":
            return await search_flight_itineraries(name, arguments)
        case _:
            raise ValueError(f"Unknown tool: {name}")


# 注册工具列表和分发函数
mcp._mcp_server.list_tools()(list_tools)
mcp._mcp_server.call_tool()(dispatch)

if __name__ == "__main__":
    mcp.run(transport='stdio')